#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
下载处理模块
负责处理蓝奏云下载页面，包括输入密码和点击下载按钮
"""

import time
import re
import traceback
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.common.exceptions import TimeoutException, NoSuchElementException
from selenium.webdriver import Keys
from selenium.webdriver.common.action_chains import ActionChains

class DownloadHandler:
    """下载处理类，处理蓝奏云下载页面相关操作"""
    
    def __init__(self, driver):
        """初始化下载处理器
        
        Args:
            driver: WebDriver对象
        """
        self.driver = driver
        self.wait = WebDriverWait(self.driver, 10)
    
    def process_download_link(self, url, password=None, auto_process=False):
        """处理下载链接，输入密码，点击下载按钮
        
        Args:
            url: 下载页面的URL
            password: 下载密码
            auto_process: 是否自动处理下载而不询问用户
            
        Returns:
            bool: 下载是否成功
        """
        try:
            print(f"\n正在处理下载链接: {url}")
            self.driver.get(url)
            time.sleep(8)  # 增加等待时间确保页面完全加载
            
            # 处理密码输入
            if password:
                if not self._handle_password_input(password):
                    print("密码输入失败，尝试继续处理...")
                time.sleep(5)  # 等待密码处理和页面刷新
            
            # 首先尝试在主框架中查找并处理下载按钮
            if self._find_and_click_download_button():
                return True
                
            # 尝试处理所有iframe
            all_iframes = self.driver.find_elements(By.TAG_NAME, "iframe")
            if all_iframes:
                print(f"找到 {len(all_iframes)} 个iframe，尝试检查每个iframe")
                
                for i, iframe in enumerate(all_iframes):
                    try:
                        print(f"切换到iframe {i+1}/{len(all_iframes)}")
                        self.driver.switch_to.frame(iframe)
                        
                        if self._find_and_click_download_button():
                            return True
                        
                        # 切回主框架
                        self.driver.switch_to.default_content()
                    except Exception as e:
                        print(f"处理iframe {i+1} 时出错: {str(e)}")
                        try:
                            self.driver.switch_to.default_content()
                        except:
                            pass
            
            # 如果所有尝试都失败，尝试从页面源码中提取下载链接
            print("尝试从页面源码提取下载链接...")
            if self._extract_download_link_from_source():
                return True
                
            # 最终失败
            print("所有下载尝试均失败")
            print(f"当前页面URL: {self.driver.current_url}")
            print("页面标题: " + self.driver.title)
            return False
                
        except Exception as e:
            print(f"处理下载链接时出错: {str(e)}")
            traceback.print_exc()
            return False
    
    def _handle_password_input(self, password):
        """处理密码输入
        
        Args:
            password: 下载密码
            
        Returns:
            bool: 是否成功输入密码
        """
        password_selectors = [
            (By.ID, "pwd"),
            (By.XPATH, "//input[@type='password']"),
            (By.CSS_SELECTOR, ".passwddiv input"),
            (By.XPATH, "//input[contains(@placeholder,'输入密码')]"),
            (By.XPATH, "//input[contains(@class,'pwd')]"),
            (By.CSS_SELECTOR, "input[type='text']"),
            (By.CSS_SELECTOR, ".d-input input")
        ]
        
        # 查找密码输入框
        pwd_input = None
        for selector_type, selector in password_selectors:
            try:
                elements = self.driver.find_elements(selector_type, selector)
                for element in elements:
                    if element.is_displayed() and element.is_enabled():
                        pwd_input = element
                        print(f"找到密码输入框: {selector_type} - {selector}")
                        break
                if pwd_input:
                    break
            except:
                continue
        
        if not pwd_input:
            print("未找到密码输入框")
            return False
            
        # 输入密码
        pwd_input.clear()
        pwd_input.send_keys(password)
        print(f"已输入密码: {password}")
        
        # 尝试查找提交按钮
        submit_selectors = [
            (By.XPATH, "//div[@class='passwddiv']/button"),
            (By.XPATH, "//button[contains(text(), '提取')]"),
            (By.XPATH, "//button[contains(text(), '确定')]"),
            (By.CSS_SELECTOR, ".passwddiv button"),
            (By.XPATH, "//button[@type='submit']"),
            (By.CSS_SELECTOR, ".d-button")
        ]
        
        submit_button = None
        for selector_type, selector in submit_selectors:
            try:
                elements = self.driver.find_elements(selector_type, selector)
                for element in elements:
                    if element.is_displayed() and element.is_enabled():
                        submit_button = element
                        print(f"找到提交按钮: {selector_type} - {selector}")
                        break
                if submit_button:
                    break
            except:
                continue
        
        if submit_button:
            try:
                submit_button.click()
                print("已点击提交密码按钮")
                return True
            except Exception as e:
                print(f"点击提交按钮失败，尝试JavaScript方式")
                try:
                    self.driver.execute_script("arguments[0].click();", submit_button)
                    print("通过JavaScript点击提交按钮成功")
                    return True
                except:
                    pass
        
        # 尝试用回车键提交
        print("未找到提交按钮或点击失败，尝试用回车键提交")
        try:
            pwd_input.send_keys(Keys.RETURN)
            print("已通过回车键提交密码")
            return True
        except:
            print("回车键提交失败")
            return False
    
    def _find_and_click_download_button(self):
        """查找并处理下载按钮
        
        Returns:
            bool: 是否成功处理下载
        """
        # 1. 首先尝试直接获取downajax中的下载链接
        try:
            downajax_div = self.driver.find_element(By.ID, "downajax")
            if downajax_div:
                print("找到ID为downajax的div元素")
                a_tags = downajax_div.find_elements(By.TAG_NAME, "a")
                for a_tag in a_tags:
                    href = a_tag.get_attribute("href")
                    if href and href.strip() and not href.startswith("javascript"):
                        print(f"找到下载链接: {href}")
                        if self._open_download_link(href):
                            return True
        except Exception as e:
            print(f"从downajax获取下载链接失败: {str(e)}")
        
        # 2. 查找具有特定ID的下载元素
        id_selectors = ["downajax", "download", "downloadfile"]
        for id_selector in id_selectors:
            try:
                element = self.driver.find_element(By.ID, id_selector)
                if element.is_displayed():
                    # 如果是a标签，尝试获取href
                    if element.tag_name == "a":
                        href = element.get_attribute("href")
                        if href and not href.startswith("javascript") and not href.startswith("#"):
                            print(f"找到ID为{id_selector}的下载链接: {href}")
                            if self._open_download_link(href):
                                return True
                    
                    # 尝试点击元素
                    print(f"尝试点击ID为{id_selector}的元素")
                    if self._try_click_element(element):
                        return True
            except:
                continue
        
        # 3. 查找包含"下载"文本的元素
        try:
            download_elements = self.driver.find_elements(By.XPATH, 
                "//*[contains(text(), '下载') or contains(text(), 'download') or contains(text(), 'Download')]")
            
            for element in download_elements:
                if not element.is_displayed() or not element.is_enabled():
                    continue
                    
                # 如果是a标签，尝试获取href
                if element.tag_name == "a":
                    href = element.get_attribute("href")
                    if href and not href.startswith("javascript") and not href.startswith("#"):
                        print(f"找到文本包含'下载'的链接: {href}")
                        if self._open_download_link(href):
                            return True
                
                # 尝试点击元素
                print("尝试点击包含'下载'文本的元素")
                if self._try_click_element(element):
                    return True
        except Exception as e:
            print(f"查找包含'下载'文本的元素失败: {str(e)}")
        
        # 4. 查找所有a标签，筛选可能的下载链接
        try:
            a_tags = self.driver.find_elements(By.TAG_NAME, "a")
            for a_tag in a_tags:
                if not a_tag.is_displayed() or not a_tag.is_enabled():
                    continue
                    
                href = a_tag.get_attribute("href")
                if not href or href.startswith("javascript") or href.startswith("#"):
                    continue
                
                # 评估这个链接是否可能是下载链接
                score = 0
                if "download" in href.lower() or "file" in href.lower():
                    score += 5
                
                text = a_tag.text.lower()
                if "下载" in text or "download" in text:
                    score += 10
                
                if score > 0:
                    print(f"找到可能的下载链接 (分数: {score}): {href}")
                    if self._open_download_link(href):
                        return True
        except Exception as e:
            print(f"查找可能的下载链接失败: {str(e)}")
        
        return False
    
    def _try_click_element(self, element):
        """尝试多种方式点击元素
        
        Args:
            element: 要点击的元素
            
        Returns:
            bool: 是否成功点击
        """
        # 尝试使用JavaScript点击
        try:
            self.driver.execute_script("arguments[0].scrollIntoView({block: 'center'});", element)
            time.sleep(1)
            self.driver.execute_script("arguments[0].click();", element)
            print("已通过JavaScript点击元素")
            time.sleep(3)
            return True
        except Exception as e:
            print(f"JavaScript点击失败: {str(e)}")
        
        # 尝试直接点击
        try:
            element.click()
            print("已直接点击元素")
            time.sleep(3)
            return True
        except Exception as e:
            print(f"直接点击失败: {str(e)}")
        
        # 尝试使用动作链点击
        try:
            actions = ActionChains(self.driver)
            actions.move_to_element(element).click().perform()
            print("已通过动作链点击元素")
            time.sleep(3)
            return True
        except Exception as e:
            print(f"动作链点击失败: {str(e)}")
            
        return False
    
    def _open_download_link(self, href):
        """打开下载链接
        
        Args:
            href: 下载链接
            
        Returns:
            bool: 是否成功打开
        """
        try:
            print(f"尝试直接打开下载链接: {href}")
            self.driver.execute_script(f"window.location.href='{href}';")
            time.sleep(5)
            print("已通过直接导航打开下载链接")
            return True
        except Exception as e:
            print(f"直接导航打开链接失败: {str(e)}")
            
            try:
                print("尝试使用get方法打开链接...")
                self.driver.get(href)
                time.sleep(5)
                print("已通过get方法打开下载链接")
                return True
            except Exception as e2:
                print(f"get方法打开链接失败: {str(e2)}")
                
            return False
    
    def _extract_download_link_from_source(self):
        """从页面源码中提取下载链接
        
        Returns:
            bool: 是否成功提取并打开链接
        """
        try:
            page_source = self.driver.page_source.lower()
            
            # 查找可能的下载链接
            download_links = re.findall(r'href=[\'"]([^\'"]*download[^\'"]*)[\'"]', page_source)
            download_links.extend(re.findall(r'href=[\'"]([^\'"]*file[^\'"]*)[\'"]', page_source))
            
            if download_links:
                for link in download_links:
                    if link and not link.startswith("javascript") and not link.startswith("#"):
                        print(f"从页面源码找到下载链接: {link}")
                        if self._open_download_link(link):
                            return True
                            
            return False
        except Exception as e:
            print(f"从源码提取下载链接失败: {str(e)}")
            return False 